# CNN用于多参数估计
import math
import torch
import torch.nn as nn
import torch.nn.functional as F

class MpeCnnModelV2(nn.Module):
    def __init__(self):
        super().__init__()
        # 第1层：256*256*8 -> 128*128*16
        self.gen_l1()
        # 第2层：128*128*16 -> 64*64*32
        self.gen_l2()
        # 第3层：64*64*32 -> 32*32*64
        self.gen_l3()
        # 第4层：32*32*64 -> 16*16*128
        self.gen_l4()
        # 第5层：16*16*128 -> 8*8*256
        self.gen_l5()
        # 第6层：8*8*256 -> 4*4*512
        self.gen_l6()
        # 第7层：4*4*512 -> 2*2*1024
        self.gen_l7()
        # 第8层：2*2*1024 -> 4096
        self.flatten = nn.Flatten()
        # 第9层：4096 -> 512
        self.fc1 = nn.Linear(2*2*1024, 512)
        # 第10层：512 -> 12
        self.fc2 = nn.Linear(512, 12)

    def forward(self, x):
        a1_real = self.l1_mp(self.l1_relu(self.l1_bn(self.l1_conv2d(torch.real(x).float()))))
        a1_imag = self.l1_mp(self.l1_relu(self.l1_bn(self.l1_conv2d(torch.imag(x).float()))))
        a2_real = self.l2_mp(self.l2_relu(self.l2_bn(self.l2_conv2d(a1_real))))
        a2_imag = self.l2_mp(self.l2_relu(self.l2_bn(self.l2_conv2d(a1_imag))))
        a3_real = self.l3_mp(self.l3_relu(self.l3_bn(self.l3_conv2d(a2_real))))
        a3_imag = self.l3_mp(self.l3_relu(self.l3_bn(self.l3_conv2d(a2_imag))))        
        a4_real = self.l4_mp(self.l4_relu(self.l4_bn(self.l4_conv2d(a3_real))))
        a4_imag = self.l4_mp(self.l4_relu(self.l4_bn(self.l4_conv2d(a3_imag))))        
        a5_real = self.l5_mp(self.l5_relu(self.l5_bn(self.l5_conv2d(a4_real))))
        a5_imag = self.l5_mp(self.l5_relu(self.l5_bn(self.l5_conv2d(a4_imag))))        
        a6_real = self.l6_mp(self.l6_relu(self.l6_bn(self.l6_conv2d(a5_real))))
        a6_imag = self.l6_mp(self.l6_relu(self.l6_bn(self.l6_conv2d(a5_imag))))       
        a7_real = self.l7_mp(self.l7_relu(self.l7_bn(self.l7_conv2d(a6_real))))
        a7_imag = self.l7_mp(self.l7_relu(self.l7_bn(self.l7_conv2d(a6_imag))))
        a8_real = self.flatten(a7_real)
        a8_imag = self.flatten(a7_imag)
        a9_real = F.relu(self.fc1(a8_real))
        a9_imag = F.relu(self.fc1(a8_imag))
        a10_real = F.relu(self.fc2(a9_real))
        a10_imag = F.relu(self.fc2(a9_imag))
        return a10_real + a10_imag

    def gen_l1(self):
        # 第1层：256*256*8 -> 128*128*16
        in_channels = 8
        out_channels = 16
        kernel = (31, 31)
        stride = (1,1)
        dilation = 1
        self.l1_conv2d = nn.Conv2d(in_channels,out_channels,kernel,stride=stride, padding='same',dilation=dilation)
        self.l1_bn = nn.BatchNorm2d(out_channels)
        self.l1_relu = nn.ReLU(True)
        mp_size, mp_stride = 2, 2
        self.l1_mp = nn.MaxPool2d(mp_size, mp_stride, padding=0)

    def gen_l2(self):
        # 第2层：128*128*16 -> 64*64*32
        in_channels = 16
        out_channels = 32
        kernel = (15, 15)
        stride = (1, 1)
        dilation = 1
        self.l2_conv2d = nn.Conv2d(in_channels,out_channels,kernel,stride=stride, padding='same',dilation=dilation)
        self.l2_bn = nn.BatchNorm2d(out_channels)
        self.l2_relu = nn.ReLU(True)
        mp_size, mp_stride = 2, 2
        self.l2_mp = nn.MaxPool2d(mp_size, mp_stride, padding=0)

    def gen_l3(self):
        # 第3层：64*64*32 -> 32*32*64
        in_channels = 32
        out_channels = 64
        kernel = (7, 7)
        stride = (1, 1)
        dilation = 1
        self.l3_conv2d = nn.Conv2d(in_channels,out_channels,kernel,stride=stride, padding='same',dilation=dilation)
        self.l3_bn = nn.BatchNorm2d(out_channels)
        self.l3_relu = nn.ReLU(True)
        mp_size, mp_stride = 2, 2
        self.l3_mp = nn.MaxPool2d(mp_size, mp_stride, padding=0)

    def gen_l4(self):
        # 第4层：32*32*64 -> 16*16*128
        in_channels = 64
        out_channels = 128
        kernel = (5, 5)
        stride = (1, 1)
        dilation = 1
        self.l4_conv2d = nn.Conv2d(in_channels,out_channels,kernel,stride=stride, padding='same',dilation=dilation)
        self.l4_bn = nn.BatchNorm2d(out_channels)
        self.l4_relu = nn.ReLU(True)
        mp_size, mp_stride = 2, 2
        self.l4_mp = nn.MaxPool2d(mp_size, mp_stride, padding=0)

    def gen_l5(self):
        # 第5层：16*16*128 -> 8*8*256
        in_channels = 128
        out_channels = 256
        kernel = (3, 3)
        stride = (1, 1)
        dilation = 1
        self.l5_conv2d = nn.Conv2d(in_channels,out_channels,kernel,stride=stride, padding='same',dilation=dilation)
        self.l5_bn = nn.BatchNorm2d(out_channels)
        self.l5_relu = nn.ReLU(True)
        mp_size, mp_stride = 2, 2
        self.l5_mp = nn.MaxPool2d(mp_size, mp_stride, padding=0)

    def gen_l6(self):
        # 第6层：8*8*256 -> 4*4*512
        in_channels = 256
        out_channels = 512
        kernel = (3, 3)
        stride = (1, 1)
        dilation = 1
        self.l6_conv2d = nn.Conv2d(in_channels,out_channels,kernel,stride=stride, padding='same',dilation=dilation)
        self.l6_bn = nn.BatchNorm2d(out_channels)
        self.l6_relu = nn.ReLU(True)
        mp_size, mp_stride = 2, 2
        self.l6_mp = nn.MaxPool2d(mp_size, mp_stride, padding=0)

    def gen_l7(self):
        # 第7层：4*4*512 -> 2*2*1024
        in_channels = 512
        out_channels = 1024
        kernel = (3, 3)
        stride = (1, 1)
        dilation = 1
        self.l7_conv2d = nn.Conv2d(in_channels,out_channels,kernel,stride=stride, padding='same',dilation=dilation)
        self.l7_bn = nn.BatchNorm2d(out_channels)
        self.l7_relu = nn.ReLU(True)
        mp_size, mp_stride = 2, 2
        self.l7_mp = nn.MaxPool2d(mp_size, mp_stride, padding=0)